use ode_solvers::*;

pub type State = Vector1<f64>;
pub type Params = Vec<f64>;

pub struct MalthusianGrowth {
    params: Params,
}

impl MalthusianGrowth {
    pub fn new(params: Option<Params>) -> Self {
        Self {
            params: params.unwrap_or(vec![1.0f64]),
        }
    }
}

impl System<State> for MalthusianGrowth {
    fn system(&self, _t: f64, y: &State, dy: &mut State) {
        dy[0] = self.params[0] * y[0];
    }
}

#[derive(Clone, Debug)]
pub struct LogisicGrowth {
    params: Params,
}

impl LogisicGrowth {
    pub fn new(params: Option<Params>) -> Self {
        Self {
            params: params.unwrap_or(vec![1.0f64, 100.0f64]),
        }
    }
}

impl System<State> for LogisicGrowth {
    fn system(&self, _t: f64, y: &State, dy: &mut State) {
        dy[0] = self.params[0] * y[0] * (1.0f64 - y[0] / self.params[1]);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_malthusian_growth() {
        let ic = State::new(1.0f64);
        let params = None;
        let malth = MalthusianGrowth::new(params);
        let tol = 1e-12;
        let mut integrator = Dopri5::new(malth, 0.0, 10.0, 0.1, ic, tol, tol);
        let result = integrator.integrate();
        match result {
            Ok(_) => {
                let t = integrator.x_out();
                let actual = integrator.y_out();
                let expected: Vec<f64> = t.iter().map(|x| x.exp()).collect();
                let mut err = 0.0f64;
                for (ix, val) in actual.iter().enumerate() {
                    err += (val[0] - expected[ix]).abs();
                }
                println!("Converged to within {err}");
                assert!(err <= 1e-6);
            }
            Err(x) => {
                println!("Integration failed with: {:?}", x);
                assert!(false);
            }
        }
    }

    #[test]
    fn test_logistic_growth() {
        let ic = State::new(1.0f64);
        let params = None;
        let logi = LogisicGrowth::new(params);
        let r = logi.params[0].clone();
        let k = logi.params[1].clone();
        let tol = 1e-12;
        let mut integrator = Dopri5::new(logi, 0.0, 10.0, 0.1, ic, tol, tol);
        let result = integrator.integrate();
        match result {
            Ok(_) => {
                let t = integrator.x_out();
                let actual = integrator.y_out();
                let expected: Vec<f64> = t
                    .iter()
                    .map(|x| {
                        let e_rt = (x * r).exp();
                        (ic[0] * k * e_rt) / (k + ic[0] * (e_rt - 1.0f64))
                    })
                    .collect();
                let mut err = 0.0f64;
                for (ix, val) in actual.iter().enumerate() {
                    err += (val[0] - expected[ix]).abs();
                }
                println!("Converged to within {err}");
                assert!(err <= 1e-6);
                assert!(true);
            }
            Err(x) => {
                println!("Integration failed with: {:?}", x);
                assert!(false);
            }
        }
    }
}
